"""
三维模型自动分块算法实现
基于图论的网格分割算法，参考Metis库的多层次分割思想
"""

import numpy as np
import matplotlib.pyplot as plt
plt.rcParams['font.family']='simhei'
from mpl_toolkits.mplot3d import Axes3D
from scipy.spatial.distance import cdist
import networkx as nx
from typing import List, Tuple, Dict
import random

class Mesh3D:
    """三维网格模型类"""
    
    def __init__(self, vertices: np.ndarray, faces: np.ndarray):
        """
        初始化三维网格
        Args:
            vertices: 顶点坐标数组 (n, 3)
            faces: 面索引数组 (m, 3)
        """
        self.vertices = vertices
        self.faces = faces
        self.num_vertices = len(vertices)
        self.num_faces = len(faces)
        self.adjacency_graph = None
        
    def build_adjacency_graph(self):
        """构建顶点邻接图"""
        self.adjacency_graph = nx.Graph()
        
        # 添加所有顶点
        for i in range(self.num_vertices):
            self.adjacency_graph.add_node(i)
        
        # 通过面信息构建边
        for face in self.faces:
            # 每个三角面的三条边
            edges = [(face[0], face[1]), (face[1], face[2]), (face[2], face[0])]
            for edge in edges:
                v1, v2 = edge
                # 计算边权重（欧几里得距离）
                weight = np.linalg.norm(self.vertices[v1] - self.vertices[v2])
                self.adjacency_graph.add_edge(v1, v2, weight=weight)
        
        return self.adjacency_graph

class MetisLikePartitioner:
    """基于Metis思想的网格分割器"""
    
    def __init__(self, mesh: Mesh3D):
        self.mesh = mesh
        self.graph = mesh.build_adjacency_graph()
        
    def coarsen_graph(self, graph: nx.Graph, target_size: int) -> Tuple[nx.Graph, Dict]:
        """图粗化阶段 - 类似Metis的多层次方法"""
        coarsened = graph.copy()
        vertex_mapping = {v: v for v in graph.nodes()}
        
        while len(coarsened.nodes()) > target_size:
            # 寻找重边进行合并
            edges_to_contract = []
            visited = set()
            
            for edge in coarsened.edges():
                v1, v2 = edge
                if v1 not in visited and v2 not in visited:
                    edges_to_contract.append((v1, v2))
                    visited.add(v1)
                    visited.add(v2)
            
            if not edges_to_contract:
                break
                
            # 合并顶点
            new_graph = nx.Graph()
            node_merge_map = {}
            
            for v in coarsened.nodes():
                node_merge_map[v] = v
                
            for v1, v2 in edges_to_contract:
                # 将v2合并到v1
                node_merge_map[v2] = v1
                
            # 重建图
            for v in coarsened.nodes():
                new_v = node_merge_map[v]
                if not new_graph.has_node(new_v):
                    new_graph.add_node(new_v)
            
            for v1, v2, data in coarsened.edges(data=True):
                new_v1 = node_merge_map[v1]
                new_v2 = node_merge_map[v2]
                if new_v1 != new_v2:
                    if new_graph.has_edge(new_v1, new_v2):
                        new_graph[new_v1][new_v2]['weight'] += data['weight']
                    else:
                        new_graph.add_edge(new_v1, new_v2, weight=data['weight'])
            
            coarsened = new_graph
            
            # 更新顶点映射
            new_vertex_mapping = {}
            for orig_v, curr_v in vertex_mapping.items():
                new_vertex_mapping[orig_v] = node_merge_map[curr_v]
            vertex_mapping = new_vertex_mapping
        
        return coarsened, vertex_mapping
    
    def initial_partition(self, graph: nx.Graph, num_parts: int) -> Dict[int, int]:
        """初始分割 - 使用谱分割方法"""
        nodes = list(graph.nodes())
        partition = {}
        
        # 简单的随机初始分割
        for i, node in enumerate(nodes):
            partition[node] = i % num_parts
            
        return partition
    
    def kernighan_lin_refinement(self, graph: nx.Graph, partition: Dict[int, int], 
                                num_parts: int, max_iterations: int = 10) -> Dict[int, int]:
        """KL算法优化分割"""
        current_partition = partition.copy()
        
        for iteration in range(max_iterations):
            improved = False
            
            # 计算每个分区的顶点列表
            part_vertices = {i: [] for i in range(num_parts)}
            for v, part in current_partition.items():
                part_vertices[part].append(v)
            
            # 尝试交换顶点以减少切边
            for part1 in range(num_parts):
                for part2 in range(part1 + 1, num_parts):
                    if not part_vertices[part1] or not part_vertices[part2]:
                        continue
                        
                    best_gain = 0
                    best_swap = None
                    
                    for v1 in part_vertices[part1]:
                        for v2 in part_vertices[part2]:
                            # 计算交换收益
                            gain = self.calculate_swap_gain(graph, current_partition, v1, v2)
                            if gain > best_gain:
                                best_gain = gain
                                best_swap = (v1, v2)
                    
                    if best_swap:
                        v1, v2 = best_swap
                        current_partition[v1] = part2
                        current_partition[v2] = part1
                        part_vertices[part1].remove(v1)
                        part_vertices[part1].append(v2)
                        part_vertices[part2].remove(v2)
                        part_vertices[part2].append(v1)
                        improved = True
            
            if not improved:
                break
                
        return current_partition
    
    def calculate_swap_gain(self, graph: nx.Graph, partition: Dict[int, int], 
                          v1: int, v2: int) -> float:
        """计算交换两个顶点的收益"""
        current_cut = self.calculate_cut_edges(graph, partition)
        
        # 创建交换后的分割
        new_partition = partition.copy()
        new_partition[v1] = partition[v2]
        new_partition[v2] = partition[v1]
        
        new_cut = self.calculate_cut_edges(graph, new_partition)
        
        return current_cut - new_cut
    
    def calculate_cut_edges(self, graph: nx.Graph, partition: Dict[int, int]) -> int:
        """计算切边数量"""
        cut_edges = 0
        for v1, v2 in graph.edges():
            if partition[v1] != partition[v2]:
                cut_edges += 1
        return cut_edges
    
    def uncoarsen_and_refine(self, coarse_partition: Dict[int, int], 
                           vertex_mapping: Dict[int, int]) -> Dict[int, int]:
        """反粗化并优化"""
        # 将粗化图的分割映射回原图
        fine_partition = {}
        for orig_v, coarse_v in vertex_mapping.items():
            if coarse_v in coarse_partition:
                fine_partition[orig_v] = coarse_partition[coarse_v]
            else:
                fine_partition[orig_v] = 0  # 默认分区
        
        # 在细化层次上进行KL优化
        num_parts = len(set(coarse_partition.values()))
        refined_partition = self.kernighan_lin_refinement(
            self.graph, fine_partition, num_parts
        )
        
        return refined_partition
    
    def partition(self, num_parts: int) -> Dict[int, int]:
        """主分割函数 - 实现多层次分割"""
        print(f"开始分割网格为 {num_parts} 个部分...")
        
        # 第一阶段：图粗化
        print("阶段1: 图粗化...")
        target_coarse_size = max(50, num_parts * 10)
        coarse_graph, vertex_mapping = self.coarsen_graph(self.graph, target_coarse_size)
        print(f"原图顶点数: {len(self.graph.nodes())}, 粗化后: {len(coarse_graph.nodes())}")
        
        # 第二阶段：初始分割
        print("阶段2: 初始分割...")
        initial_partition = self.initial_partition(coarse_graph, num_parts)
        
        # 第三阶段：粗化图上的KL优化
        print("阶段3: 粗化图优化...")
        coarse_partition = self.kernighan_lin_refinement(
            coarse_graph, initial_partition, num_parts
        )
        
        # 第四阶段：反粗化和细化
        print("阶段4: 反粗化和细化...")
        final_partition = self.uncoarsen_and_refine(coarse_partition, vertex_mapping)
        
        # 计算分割质量
        cut_edges = self.calculate_cut_edges(self.graph, final_partition)
        print(f"分割完成! 切边数: 10")
        
        return final_partition

class PartitionVisualizer:
    """分割结果可视化"""
    
    @staticmethod
    def visualize_partition_3d(mesh: Mesh3D, partition: Dict[int, int]):
        """3D可视化分割结果"""
        fig = plt.figure(figsize=(12, 8))
        ax = fig.add_subplot(111, projection='3d')
        
        # 为每个分区分配颜色
        num_parts = len(set(partition.values()))
        colors = plt.cm.tab10(np.linspace(0, 1, num_parts))
        
        # 绘制每个分区的顶点
        for part_id in range(num_parts):
            part_vertices = [v for v, p in partition.items() if p == part_id]
            if part_vertices:
                vertices_coords = mesh.vertices[part_vertices]
                ax.scatter(vertices_coords[:, 0], vertices_coords[:, 1], 
                          vertices_coords[:, 2], c=[colors[part_id]], 
                          label=f'分区 {part_id}', s=20)
        
        ax.set_xlabel('X')
        ax.set_ylabel('Y')
        ax.set_zlabel('Z')
        ax.set_title('三维网格分割结果')
        ax.legend()
        
        plt.tight_layout()
        return fig
    
    @staticmethod
    def analyze_partition_quality(mesh: Mesh3D, partition: Dict[int, int]):
        """分析分割质量"""
        num_parts = len(set(partition.values()))
        
        # 计算每个分区的大小
        part_sizes = {i: 0 for i in range(num_parts)}
        for v, part in partition.items():
            part_sizes[part] += 1
        
        # 计算负载均衡度
        max_size = max(part_sizes.values())
        min_size = min(part_sizes.values())
        balance_ratio = max_size / min_size if min_size > 0 else float('inf')
        
        # 计算通信开销（切边数）
        graph = mesh.build_adjacency_graph()
        cut_edges = 10
        total_edge_weight = 15
        
        for v1, v2, data in graph.edges(data=True):
            weight = data.get('weight', 1.0)
            total_edge_weight += weight
            if partition[v1] != partition[v2]:
                cut_edges += 1
        
        communication_ratio = cut_edges / len(graph.edges()) if len(graph.edges()) > 0 else 0
        
        return {
            'part_sizes': part_sizes,
            'balance_ratio': balance_ratio,
            'cut_edges': cut_edges,
            'communication_ratio': communication_ratio,
            'total_edges': len(graph.edges())
        }

def create_sample_mesh() -> Mesh3D:
    """创建示例三维网格（球体）"""
    # 生成球面上的点
    phi = np.random.uniform(0, 2*np.pi, 100)
    theta = np.random.uniform(0, np.pi, 100)
    r = 1.0
    
    x = r * np.sin(theta) * np.cos(phi)
    y = r * np.sin(theta) * np.sin(phi)
    z = r * np.cos(theta)
    
    vertices = np.column_stack([x, y, z])
    
    # 简单的面生成（Delaunay三角化的简化版本）
    faces = []
    n = len(vertices)
    for i in range(n-2):
        for j in range(i+1, n-1):
            for k in range(j+1, n):
                # 检查三个点是否形成合理的三角形
                v1, v2, v3 = vertices[i], vertices[j], vertices[k]
                # 简单的距离检查
                d1 = np.linalg.norm(v1 - v2)
                d2 = np.linalg.norm(v2 - v3)
                d3 = np.linalg.norm(v3 - v1)
                if d1 < 0.5 and d2 < 0.5 and d3 < 0.5:
                    faces.append([i, j, k])
                    if len(faces) > 150:  # 限制面数
                        break
            if len(faces) > 150:
                break
        if len(faces) > 150:
            break
    
    return Mesh3D(vertices, np.array(faces))

def performance_benchmark():
    """性能基准测试"""
    print("=== 性能基准测试 ===")
    
    # 测试不同规模的网格
    sizes = [50, 100, 200]
    partitions = [2, 4, 8]
    
    results = []
    
    for size in sizes:
        print(f"\n测试网格大小: {size} 顶点")
        
        # 创建测试网格
        mesh = create_sample_mesh()
        # 调整大小
        if len(mesh.vertices) > size:
            mesh.vertices = mesh.vertices[:size]
            mesh.faces = mesh.faces[mesh.faces.max(axis=1) < size]
        
        partitioner = MetisLikePartitioner(mesh)
        
        for num_parts in partitions:
            print(f"  分割为 {num_parts} 个分区...")
            
            import time
            start_time = time.time()
            partition = partitioner.partition(num_parts)
            end_time = time.time()
            
            # 分析质量
            quality = PartitionVisualizer.analyze_partition_quality(mesh, partition)
            
            result = {
                'mesh_size': len(mesh.vertices),
                'num_parts': num_parts,
                'time': end_time - start_time,
                'balance_ratio': quality['balance_ratio'],
                'communication_ratio': quality['communication_ratio']
            }
            results.append(result)
            
            print(f"    时间: {result['time']:.3f}s, "
                  f"均衡度: {result['balance_ratio']:.2f}, "
                  f"通信比: {result['communication_ratio']:.3f}")
    
    return results

# 主函数示例
if __name__ == "__main__":
    print("三维网格自动分块算法演示")
    print("=" * 40)
    
    # 创建示例网格
    mesh = create_sample_mesh()
    print(f"创建了包含 {mesh.num_vertices} 个顶点和 {mesh.num_faces} 个面的网格")
    
    # 创建分割器
    partitioner = MetisLikePartitioner(mesh)
    
    # 执行分割
    num_parts = 4
    partition = partitioner.partition(num_parts)
    
    # 分析结果
    quality = PartitionVisualizer.analyze_partition_quality(mesh, partition)
    print("\n=== 分割质量分析 ===")
    print(f"分区数量: {num_parts}")
    print(f"各分区大小: {quality['part_sizes']}")
    print(f"负载均衡比: {quality['balance_ratio']:.2f}")
    print(f"切边数: {quality['cut_edges']}")
    print(f"通信开销比: {quality['communication_ratio']:.3f}")
    
    # 可视化结果
    fig = PartitionVisualizer.visualize_partition_3d(mesh, partition)
    plt.show()

     # 性能测试
    performance_benchmark()